import onnx
from onnx_tf.backend import prepare
import tensorflow as tf
import numpy as np
from PIL import Image


def preprocess_image(image_path):
    """
    将任意分辨率的图像加载、灰度化、缩放到 28x28 并归一化，以符合 MNIST 模型输入要求。
    """
    img = Image.open(image_path).convert("L")
    img = img.resize((28, 28))  # 缩放到 28x28
    img_array = np.array(img, dtype=np.float32) / 255.0  # 归一化到 [0, 1]
    img_array = img_array.reshape(1, 28, 28, 1)  # 调整为 (1, 28, 28, 1)
    return img_array


def predict(image_path):
    """
    使用 TensorFlow 读取 ONNX 模型并进行推理。
    """
    # 加载 ONNX 模型并转换为 TensorFlow 模型
    onnx_model = onnx.load("mnist_model.onnx")
    tf_model = prepare(onnx_model)  # 使用 onnx-tf 转换

    # 预处理输入图像
    input_data = preprocess_image(image_path)

    # 进行预测
    prediction = tf_model.run(input_data)

    # 获取预测类别
    predicted_class = np.argmax(prediction[0])
    print(f"预测类别: {predicted_class}")


# 使用任意分辨率图像的路径
image_path = "your_image_path.png"  # 替换为图片路径
predict(image_path)
